package textprocessing.nlpanalysis.coreference;


import edu.stanford.nlp.coref.CorefCoreAnnotations;
import edu.stanford.nlp.coref.data.CorefChain;
import edu.stanford.nlp.coref.data.Mention;
import edu.stanford.nlp.ling.CoreAnnotations;
import edu.stanford.nlp.ling.CoreLabel;
import edu.stanford.nlp.pipeline.Annotation;
import edu.stanford.nlp.pipeline.StanfordCoreNLP;
import edu.stanford.nlp.util.CoreMap;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Properties;

/**
 * Created by julianzliu on 4/19/2017.
 */
public class CoReferencer {


    private Annotation document;

    public CoReferencer(Annotation document){
        this.document = document;
    }


    public void printCoref(){
        for ( CorefChain cc : this.document.get(CorefCoreAnnotations.CorefChainAnnotation.class).values() ) {
            System.out.println("\t" + cc);
        }

        for (CoreMap sentence : document.get(CoreAnnotations.SentencesAnnotation.class)) {
            System.out.println("---");
            System.out.println("mentions");
            for (Mention m : sentence.get(CorefCoreAnnotations.CorefMentionsAnnotation.class)) {
                System.out.println("\t" + m);
            }
        }
    }



    public static String resolve_coref(Annotation doc){



        Map<Integer, CorefChain> corefs = doc.get(CorefCoreAnnotations.CorefChainAnnotation.class);
        List<CoreMap> sentences = doc.get(CoreAnnotations.SentencesAnnotation.class);


        List<String> resolved = new ArrayList<String>();

        for (CoreMap sentence : sentences) {

            List<CoreLabel> tokens = sentence.get(CoreAnnotations.TokensAnnotation.class);

            for (CoreLabel token : tokens) {

                System.out.println();
                System.out.println();
                System.out.println("-----------------------------------------------------------------");

                Integer corefClustId= token.get(CorefCoreAnnotations.CorefClusterIdAnnotation.class);
                System.out.println(token.word() +  " --> corefClusterID = " + corefClustId);


                CorefChain chain = corefs.get(corefClustId);
                System.out.println("matched chain = " + chain);


                if(chain==null){
                    resolved.add(token.word());
                }else{

                    CorefChain.CorefMention reprMent = chain.getRepresentativeMention();
                    System.out.println(reprMent);

                    int sentINdx = reprMent.sentNum -1;
                    CoreMap corefSentence = sentences.get(sentINdx);
                    List<CoreLabel> corefSentenceTokens = corefSentence.get(CoreAnnotations.TokensAnnotation.class);


                    String newwords = "";

                    System.out.println("token.sentIndex = " + token.sentIndex() + " , rep.sentNum = " + sentINdx);
                    System.out.println("token.index = " + token.index() + " , rep.start = " + reprMent.startIndex + " , rep.end = " + reprMent.endIndex);
                    if(token.sentIndex() != sentINdx || token.index() < reprMent.startIndex || token.index() > reprMent.endIndex ){
                        for(int i = reprMent.startIndex; i<reprMent.endIndex; i++){
                            CoreLabel matchedLabel = corefSentenceTokens.get(i-1); //resolved.add(tokens.get(i).word());
                            resolved.add(matchedLabel.word());

                            newwords+=matchedLabel.word()+" ";
                        }
                    }
                    else{
                        resolved.add(token.word());
                    }


                    System.out.println("converting " + token.word() + " to " + newwords);
                }

                System.out.println("-----------------------------------------------------------------");

            }

        }


        String resolvedStr ="";
        System.out.println();
        for (String str : resolved) {
            resolvedStr+=str+" ";
        }
        System.out.println(resolvedStr);

        return resolvedStr;


    }


    /***************************************************************
     *
     * @param args
     */
    public static void main(String[] args){
        Properties props = new Properties();
        // coref vs dcoref
        //props.setProperty("annotators", "tokenize,ssplit,pos,lemma,ner,mention,coref");
        props.setProperty("annotators", "tokenize,ssplit,pos,lemma,ner,parse,mention,coref");
        StanfordCoreNLP pipeline = new StanfordCoreNLP(props);

        Annotation document = new Annotation("Barack Obama was born in Hawaii. He is the president of United States.");
        pipeline.annotate(document);

        CoReferencer coReferencer = new CoReferencer(document);

        coReferencer.printCoref();




    }




}
